# -*- coding: utf-8 -*-
# @Time    : 2020/6/21 下午3:09
# @Author  : caotian
# @FileName: alexnetmyopiatrain.py
# @Software: PyCharm

import os
import sys

curpath = os.path.abspath(os.curdir)
sys.path.append(curpath)
import lenetmyopiadataloader as lmd
import myopiamodel as mm
import alexnetmyopiamodel as amm
import vggmyopiamodel as vgg
import googlelenetmyopiamodel as googlelenet

import random
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D, Pool2D, Linear
import numpy as np

traindatadir = './data/PALM-Training400/PALM-Training400'
validationdatadir = './data/PALM-Validation400'
csvfile = './data/PALM-Validation-GT/labels.csv'


def train(model):
    with fluid.dygraph.guard():
        print('start training ...')
        model.train()
        epoch_num = 5
        opt = fluid.optimizer.Momentum(learning_rate=0.001, momentum=0.9, parameter_list=model.parameters())
        train_loader = lmd.data_loader(traindatadir, batch_size=10, mode='train')
        valid_loader = lmd.valid_data_loader(validationdatadir, csvfile)
        for epoch in range(epoch_num):
            for batch_id, data in enumerate(train_loader()):
                xdata, ydata = data
                img = fluid.dygraph.to_variable(xdata)
                label = fluid.dygraph.to_variable(ydata)
                logits = model(img)
                loss = fluid.layers.sigmoid_cross_entropy_with_logits(logits, label)
                avg_loss = fluid.layers.mean(loss)
                if batch_id % 10 == 0:
                    print("epoch {},batch_id {},loss {}".format(epoch, batch_id, avg_loss.numpy()))
                avg_loss.backward()
                opt.minimize(avg_loss)
                model.clear_gradients()
            model.eval()
            accuracies = []
            losses = []
            for batch_id, data in enumerate(valid_loader()):
                xdata, ydata = data
                img = fluid.dygraph.to_variable(xdata)
                label = fluid.dygraph.to_variable(ydata)
                logits = model(img)
                pred = fluid.layers.sigmoid(logits)
                loss = fluid.layers.sigmoid_cross_entropy_with_logits(pred, label)
                pred2 = pred * (-1.0) + 1.0
                pred = fluid.layers.concat([pred2, pred], axis=1)
                acc = fluid.layers.accuracy(pred, fluid.layers.cast(label, dtype='int64'))
                accuracies.append(acc.numpy())
                losses.append(loss.numpy())
            print("validation accuracy/loss:{}/{}".format(np.mean(accuracies), np.mean(losses)))
            model.train()
        fluid.save_dygraph(model.state_dict(), 'myopia-mnist')
        fluid.save_dygraph(opt.state_dict(), 'myopia-mnist')


def evaluation(model, params_file_path):
    with fluid.dygraph.guard():
        print("start evaluation...")
        model_state_dict, _ = fluid.load_dygraph(params_file_path)
        model.load_dict(model_state_dict)
        model.eval()
        eval_loader = lmd.load_data('eval')
        acc_set = []
        avg_loss_set = []
        for batch_id, data in enumerate(eval_loader()):
            xdata, ydata = data
            img = fluid.dygraph.to_variable(xdata)
            label = fluid.dygraph.to_variable(ydata)
            prediction, acc = model(img, label)
            loss = fluid.layers.cross_entropy(input=prediction, label=label)
            avg_loss = fluid.layers.mean(loss)
            acc_set.append(float(acc.numpy()))
            avg_loss_set.append(float(avg_loss.numpy()))
        acc_val_mean = np.array(acc_set).mean()
        avg_loss_val_mean = np.array(avg_loss_set).mean()
        print('loss {},acc {}'.format(avg_loss_val_mean, acc_val_mean))

import residualnetmyopiamodel as rmm
if __name__ == '__main__':
    with fluid.dygraph.guard():
        #model = mm.LeNet(num_classes=1)
        #model = amm.AlexNet(num_classes=1)
        #model=vgg.VGG()
        model=googlelenet.GoogleLeNet()
        #model=rmm.ResNet()
    train(model)

